package edu.colorado.cs.correll.cudroplet.shaperecognition;

import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.Stack;

import com.vividsolutions.jts.geom.Coordinate;
import com.vividsolutions.jts.geom.Envelope;
import com.vividsolutions.jts.index.kdtree.KdNode;
import com.vividsolutions.jts.index.kdtree.KdTree;

public abstract class Swarm {
	ShapeRecognition shapeRecognition;

	public static final int SWARM_SIZE = 1000;
	
	Robot[] robots;

	Set<Robot>[] neighbors;

	KdTree tree = new KdTree();

	Set<Subgraph> subgraphs = new HashSet<Subgraph>();

	RobotSettings robotSettings = new RobotSettings(1, 0.1f, 10, 40);

	public abstract void update();

	public void updateRobots() {
		for (int i = 0; i < robots.length; i++) {
			robots[i].update();
		}
	}

	public void updateNeighbors() {
		for (int i = 0; i < robots.length; i++) {
			if (robots[i].state == RobotState.FOUND) {
				neighbors[i].clear();
				Robot r = robots[i];
				Coordinate p = r.getCoordinate();
				Envelope env = new Envelope(new Coordinate(r.x
						- robotSettings.comm_radius,
						r.y - robotSettings.comm_radius),
				new Coordinate(r.x
						+ robotSettings.comm_radius,
						r.y + robotSettings.comm_radius));
				@SuppressWarnings("unchecked")
				List<KdNode> range = tree.query(env);

				for (KdNode kdNode: range) {
					Robot r2 = (Robot) kdNode.getData();
					Coordinate p2 = r2.getCoordinate();
					double d = p.distance(p2);
					if (d > 0
							&& d < robotSettings.comm_radius) {
						neighbors[i].add(r2);
					}
				}
			}
		}
	}

	public void calcSubgraphs() {
		subgraphs.clear();

		// set leaders to null
		for (int i = 0; i < robots.length; i++) {
			robots[i].leader = null;
		}

		for (int i = 0; i < robots.length; i++) {
			if (robots[i].state == RobotState.FOUND && robots[i].leader == null) {
				subgraphs.add(getSubgraph(robots[i]));
			}
		}
	}

	Subgraph getSubgraph(Robot root) {
		Leader leader = new Leader();
		Set<Robot> subgraphSet = new HashSet<Robot>();
		Stack<Robot> stack = new Stack<Robot>();
		// set the leader when adding to queue
		root.leader = leader;
		stack.push(root);
		subgraphSet.add(root);
		while (!stack.empty()) {
			Robot robot = stack.pop();
			Set<Robot> n = neighbors[robot.id];
			for (Robot r : n) {
				if (!subgraphSet.contains(r)) {
					r.leader = leader;
					stack.push(r);
					subgraphSet.add(r);
				}
			}
		}
		return new Subgraph(shapeRecognition, subgraphSet, leader);
	}

	public int getColor(int id) {
		return id * (255 * 255 * 255) / robots.length;
	}
}
